
#include "global_avg_pool_tiling.h"
#include "tiling/platform/platform_ascendc.h"
#include "register/op_def_registry.h"



namespace optiling {
const uint32_t BLOCK_SIZE = 32;
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

    GlobalAvgPoolTilingData tiling;
    uint32_t NUM = 12;   //
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    uint64_t ub_size; ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub_size);
    auto aivNum = ascendcPlatform.GetCoreNum();     // vector core num  1

    // 元素总量
    uint32_t total_length = context->GetInputTensor(0)->GetShapeSize();

    auto dt = context->GetInputTensor(0)->GetDataType();
    uint32_t sizeofdatatype;
    if (dt == 1) { //fp16
        sizeofdatatype = 2;
        NUM = 8;
        tiling.set_typeKey(1);
    }
    else if(dt == 0) { //float
        sizeofdatatype = 4;
        NUM = 5;
        tiling.set_typeKey(0);
    }

    uint32_t ALIGN_NUM = BLOCK_SIZE / sizeofdatatype;

    auto shape = context->GetInputTensor(0)->GetShape().GetOriginShape();
    uint32_t dim_num = shape.GetDimNum();
    uint32_t group_num = shape.GetDim(0) * shape.GetDim(1);
    uint32_t reduce_num = 1;
    for(uint32_t i = 2; i < dim_num; i++) {
        reduce_num *= shape.GetDim(i);
    }

    // if(reduce_num * sizeofdatatype % 256 == 0)
    // {
    //     context->SetTilingKey(1);   
    // }
    // else
    // {
    //     context->SetTilingKey(2);
    // }

    // uint32_t rest = (ub_size / 2 / sizeofdatatype) - group_num;
    // uint32_t packNumber = rest / (reduce_num * 2);
    // for (int i = 0; i < 20; ++i) {
    //     if ((packNumber >> i) > 1) {
    //         packNumber &= ~(1 << i);
    //     }
    // }
    // while (group_num % packNumber) {
    //     packNumber >>= 1;
    // }
    // tiling.set_packNumber(packNumber);

    // reduce_num 32向上字节对齐
    uint32_t reduce_num_Aligned = (reduce_num + ALIGN_NUM - 1) / ALIGN_NUM * ALIGN_NUM;
    // uint32_t group_num_Aligned = (group_num + ALIGN_NUM - 1) / ALIGN_NUM * ALIGN_NUM;

    // 元素总量 32Byte对齐
    // uint32_t totalLengthAligned = (total_length + ALIGN_NUM - 1) / ALIGN_NUM * ALIGN_NUM;

    // 多少个32字节 块
    uint32_t tiling_size = ((ub_size) / BLOCK_SIZE / 2) / NUM;

    // 确保block_size能256字节对齐（每次tiling的数据256字节对齐）(矢量計算單元每次讀取256字節進行運算，對齊后效率更高 )
    // tiling_size = tiling_size <=8 ? tiling_size : tiling_size / 8 * 8;
    
    // 每次tiling的数据个数 (尽可能用满UB size)
    uint32_t block_size = tiling_size * ALIGN_NUM;

    uint32_t tile_num;
    uint32_t tileLength = 0;
    uint32_t lasttileLength = 0;

    // uint32_t grouptile_num;
    // uint32_t grouptileLength = 0;
    // uint32_t grouplasttileLength = 0;

    // tiling循环次数
    tile_num = reduce_num_Aligned / block_size;
    // 
    if(reduce_num_Aligned % block_size == 0 || tile_num == 0) //reduce dim对齐之后，可以均分到各个tile，或不足一个tile
    {
        if(tile_num == 0)    //一次tiling的数据个数 大于 totalLengthAligned，也就是tile_num = 0时，tileLength=totalLengthAligned
        {
            tile_num = 1;
            tileLength = reduce_num_Aligned;
            lasttileLength = tileLength;
        }
        else //一次tiling的数据个数 小于 totalLengthAligned时，tileLength取最大值tiling_size * ALIGN_NUM
        {
            tileLength = block_size;
            lasttileLength = tileLength;
        }
    }
    else //reduce_num_Aligned对齐后不能均分到每次tiling
    {
        tile_num = tile_num + 1;
        tileLength = block_size; //tileLength 取最大值(ub)
        lasttileLength = reduce_num_Aligned % tileLength;
    }



    

    // aivNum = ((aivNum < totalLength / block_size) ? aivNum : (totalLength / block_size));
    // aivNum = aivNum >= 1 ? aivNum : 1;
    // printf("ub_size: %u \n", ub_size);
    // printf("tile_num: %u \n", tile_num);
    // printf("tileLength: %u \n", tileLength);
    // printf("lasttileLength: %u \n", lasttileLength);
    // printf("grouptile_num: %u \n", grouptile_num);
    // printf("grouptileLength: %u \n", grouptileLength);
    // printf("grouplasttileLength: %u \n", grouplasttileLength);
    // printf("reduce_num: %u \n", reduce_num);
    // printf("group_num: %u \n", group_num);
    // printf("total_length: %u \n", total_length);
    // printf("packNumber: %u \n", packNumber);
    

    context->SetBlockDim(1);    //310B4只有一个core
    tiling.set_reduceNum(reduce_num);
    tiling.set_groupNum(group_num);
    tiling.set_totalLength(total_length);
    tiling.set_tileNum(tile_num);
    tiling.set_tileLength(tileLength);
    tiling.set_lasttileLength(lasttileLength);

    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());

    size_t* currentWorkspace = context->GetWorkspaceSizes(1);
    currentWorkspace[0] = 0;

    return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class GlobalAvgPool : public OpDef {
public:
    explicit GlobalAvgPool(const char* name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend310b");

    }
};

OP_ADD(GlobalAvgPool);
}
